{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "EgiF12Hf1Dhs"
   },
   "source": [
    "This notebook provides examples to go along with the [textbook](http://manipulation.csail.mit.edu/force.html).  I recommend having both windows open, side-by-side!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "eeMrMI0-1Dhu"
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "import numpy as np\n",
    "from pydrake.all import (\n",
    "    AddMultibodyPlantSceneGraph,\n",
    "    Box,\n",
    "    DiagramBuilder,\n",
    "    LogVectorOutput,\n",
    "    MeshcatVisualizer,\n",
    "    Parser,\n",
    "    Simulator,\n",
    "    Sphere,\n",
    "    StartMeshcat,\n",
    ")\n",
    "\n",
    "from manipulation import running_as_notebook"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start the visualizer.\n",
    "meshcat = StartMeshcat()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Getting some bounce\n",
    "\n",
    "It's possible to [achieve bouncing](https://www.youtube.com/watch?v=m5UnMWihWC4&t=108s) using the default contact parameters [(models)](https://github.com/mattbev/robot-juggler/tree/main/utils/models). But adjusting those parameters can get you more bounce.  Here is a simulation of a ball being dropped from 1m. \n",
    "\n",
    "The two most relevant parameters are the contact stiffness and damping (dissipation).  Search for \"point_contact_stiffness\" and/or \"hunt_crossley_dissipation\" in the [MultibodyPlant documentation]([hunt_crossley_dissipation](https://drake.mit.edu/doxygen_cxx/classdrake_1_1multibody_1_1_multibody_plant.html)). The way to specify them in the SDF is documented [here](https://drake.mit.edu/doxygen_cxx/group__multibody__parsing.html).\n",
    "\n",
    "Note: This example uses TAMSI as the contact solver.  The SAP solver uses slightly different parameters."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def ball_drop(stiffness, dissipation):\n",
    "    sdf = f\"\"\"\n",
    "<?xml version=\"1.0\"?>\n",
    "<sdf version=\"1.7\">\n",
    "  <model name=\"bouncing_ball\">\n",
    "    <link name=\"ball\">\n",
    "      <pose>0 0 0.1 0 0 0</pose> \n",
    "      <inertial>\n",
    "        <mass>1</mass>\n",
    "        <inertia>\n",
    "          <ixx>0.001</ixx>\n",
    "          <iyy>0.001</iyy>\n",
    "          <izz>0.001</izz>\n",
    "          <ixy>0</ixy>\n",
    "          <ixz>0</ixz>\n",
    "          <iyz>0</iyz>\n",
    "        </inertia>\n",
    "      </inertial>\n",
    "      <visual name=\"visual\">\n",
    "        <geometry>\n",
    "          <sphere>\n",
    "            <radius>0.1</radius>\n",
    "          </sphere>\n",
    "        </geometry>\n",
    "        <material>\n",
    "          <diffuse>1.0 0.34 0.25 1.0</diffuse>\n",
    "        </material>\n",
    "      </visual>\n",
    "      <collision name=\"collision\">\n",
    "        <geometry>\n",
    "          <sphere>\n",
    "            <radius>0.1</radius>\n",
    "          </sphere>\n",
    "        </geometry>\n",
    "        <drake:proximity_properties>\n",
    "          <drake:point_contact_stiffness>\n",
    "            {stiffness}\n",
    "          </drake:point_contact_stiffness>\n",
    "          <drake:hunt_crossley_dissipation>\n",
    "            {dissipation}\n",
    "          </drake:hunt_crossley_dissipation>\n",
    "        </drake:proximity_properties>\n",
    "      </collision>\n",
    "    </link>\n",
    "    <joint name=\"ball_z\" type=\"prismatic\">\n",
    "      <parent>world</parent>\n",
    "      <child>ball</child>\n",
    "      <axis>\n",
    "        <xyz>0 0 1</xyz>\n",
    "        <limit>\n",
    "          <effort>0</effort>\n",
    "        </limit>\n",
    "      </axis>\n",
    "    </joint>\n",
    "    <link name=\"ground\">\n",
    "      <pose>0 0 -0.5 0 0 0</pose>\n",
    "      <visual name=\"visual\">\n",
    "        <geometry>\n",
    "          <box>\n",
    "            <size>2 2 1</size>\n",
    "          </box>\n",
    "        </geometry>\n",
    "      </visual>\n",
    "      <collision name=\"collision\">\n",
    "        <geometry>\n",
    "          <box>\n",
    "            <size>2 2 1</size>\n",
    "          </box>\n",
    "        </geometry>\n",
    "        <drake:proximity_properties>\n",
    "          <drake:point_contact_stiffness>\n",
    "            {stiffness}\n",
    "          </drake:point_contact_stiffness>\n",
    "          <drake:hunt_crossley_dissipation>\n",
    "            {dissipation}\n",
    "          </drake:hunt_crossley_dissipation>\n",
    "        </drake:proximity_properties>\n",
    "      </collision>\n",
    "    </link>\n",
    "    <joint name=\"ground_weld\" type=\"fixed\">\n",
    "      <parent>world</parent>\n",
    "      <child>ground</child>\n",
    "      <pose>0 0 0 0 0 0</pose>\n",
    "    </joint>\n",
    "  </model>\n",
    "</sdf>\n",
    "\"\"\"\n",
    "    builder = DiagramBuilder()\n",
    "    time_step = 0.001\n",
    "    plant, scene_graph = AddMultibodyPlantSceneGraph(builder, time_step)\n",
    "    Parser(plant).AddModelsFromString(sdf, \".sdf\")\n",
    "    plant.Finalize()\n",
    "\n",
    "    visualizer = MeshcatVisualizer.AddToBuilder(builder, scene_graph, meshcat)\n",
    "    logger = LogVectorOutput(plant.get_state_output_port(), builder)\n",
    "\n",
    "    diagram = builder.Build()\n",
    "    simulator = Simulator(diagram)\n",
    "\n",
    "    context = simulator.get_mutable_context()\n",
    "    plant_context = plant.GetMyMutableContextFromRoot(context)\n",
    "    plant.SetPositions(plant_context, [1])\n",
    "\n",
    "    visualizer.StartRecording()\n",
    "    simulator.AdvanceTo(3 if running_as_notebook else 0.1)\n",
    "    visualizer.PublishRecording()\n",
    "\n",
    "    log = logger.FindLog(context)\n",
    "    t = log.sample_times()\n",
    "    z = log.data()[0]\n",
    "    plt.plot(t, z)\n",
    "    first_contact = np.argmax(z < 0)\n",
    "    bounce_height = np.max(z[first_contact:])\n",
    "    print(f\"bounce height = {bounce_height}\")\n",
    "\n",
    "\n",
    "# These are the default values that MultibodyPlant picked (on Nov 16, 2022) for\n",
    "# this system.\n",
    "ball_drop(stiffness=19620, dissipation=10.096375546923044)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ball_drop(stiffness=1000, dissipation=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ball_drop(stiffness=20000, dissipation=0.1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "collapsed_sections": [],
   "name": "Robotic Manipulation - Motion Planning.ipynb",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3.10.8 64-bit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.8"
  },
  "vscode": {
   "interpreter": {
    "hash": "b0fa6594d8f4cbf19f97940f81e996739fb7646882a419484c72d19e05852a7e"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
